#include <stdlib.h>
#include <string.h>
#include <cuda_runtime.h>

#include "dvmgpu_distribution.h"

typedef struct {
    int dimension;
    LoopBounds *loopBounds;
    int threads[3];
    int resBlocks;
    IndexType *blocksAddress;
} CacheEntry;
typedef struct {
    int count;
    int beginIndex;
    int size;
    CacheEntry *entries;
} Cache;

static int add_to_cache(Cache *cache, int dimension, LoopBounds *loopBounds, int threads[3], int resBlocks, IndexType *blocksAddress) {
    int nextIndex = (cache->beginIndex + cache->count) % cache->size;
    int res = 0;
    CacheEntry *nextEntry = &cache->entries[nextIndex];
    if (cache->count == cache->size) {
        safe_free((void **)&nextEntry->loopBounds);
        assertCuda(cudaFree(nextEntry->blocksAddress));
        cache->count--;
        cache->beginIndex = (cache->beginIndex + 1) % cache->size;
        res = 1;
    }
    cache->count++;
    nextEntry->dimension = dimension;
    nextEntry->loopBounds = (LoopBounds *)malloc(sizeof(LoopBounds) * dimension);
    int i;
    for (i = 0; i < dimension; i++) {
        int j;
        for (j = 0; j < 3; j++)
            nextEntry->loopBounds[i][j] = loopBounds[i][j];
    }
    memcpy(nextEntry->threads, threads, sizeof(int) * 3);
    nextEntry->resBlocks = resBlocks;
    nextEntry->blocksAddress = blocksAddress;
    return res;
}

static int probe_cache(Cache *cache, int dimension, LoopBounds *loopBounds, int threads[3], int *pResBlocks, IndexType **pBlocksAddress) {
    int i;
    for (i = 0; i < cache->count; i++) {
        CacheEntry *entry = &cache->entries[(cache->beginIndex + cache->count - 1 - i) % cache->size];
        int flag = entry->dimension == dimension;
        int j;
        for (j = 0; flag && j < dimension; j++) {
            int k;
            for (k = 0; flag && k < 3; k++)
                flag = flag && entry->loopBounds[j][k] == loopBounds[j][k];
        }
        flag = flag && (entry->threads[0] == threads[0] && entry->threads[1] == threads[1] && entry->threads[2] == threads[2]);
        if (flag) {
            *pResBlocks = entry->resBlocks;
            *pBlocksAddress = entry->blocksAddress;
            return 1;
        }
    }
    return 0;
}

static Cache *init_cache(int size) {
    Cache *cache;
    cache = (Cache *)malloc(sizeof(Cache));
    cache->count = 0;
    cache->beginIndex = 0;
    cache->size = size;
    cache->entries = (CacheEntry *)malloc(sizeof(CacheEntry) * cache->size);
    return cache;
}

void genBlocksInfo(int dimension, LoopBounds *loopBounds, int threads[3], int *pResBlocks, IndexType **pBlocksAddress) {
    IndexType (*blocksInfo)[dimension][2];
    int resBlocks;

    int steps[dimension];
    int i;
    for (i = 0; i < dimension - 3; i++)
        steps[i] = 1;
    for (i = 0; i < 3 && i < dimension; i++)
        steps[dimension - 1 - i] = threads[i];

    resBlocks = 1;
    for (i = 0; i < dimension; i++) {
        int count = (loopBounds[i][1] - loopBounds[i][0]) / loopBounds[i][2] + 1;
        resBlocks *= (count + steps[i] - 1) / steps[i];
    }
    blocksInfo = (typeof(blocksInfo))malloc(sizeof(*blocksInfo) * resBlocks);
    int curVals[dimension], curLens[dimension];
    for (i = 0; i < dimension; i++) {
        curVals[i] = loopBounds[i][0];
        int rest = (loopBounds[i][1] - curVals[i]) / loopBounds[i][2] + 1;
        if (rest <= steps[i])
            curLens[i] = rest;
        else
            curLens[i] = steps[i];
    }
    for (i = 0; i < resBlocks; i++) {
        int j;
        for (j = 0; j < dimension; j++) {
            blocksInfo[i][j][0] = curVals[j];
            blocksInfo[i][j][1] = curVals[j] + (curLens[j] - 1) * loopBounds[j][2];
        }
// Perform step
        for (j = dimension - 1; j >= 0; j--) {
            curVals[j] += loopBounds[j][2] * curLens[j];
            int rest = (loopBounds[j][1] - curVals[j]) / loopBounds[j][2] + 1;
            if (rest <= 0)
                curVals[j] = loopBounds[j][0];
            else {
                if (rest <= steps[j])
                    curLens[j] = rest;
                else
                    curLens[j] = steps[j];
                break;
            }
        }
    }
    *pResBlocks = resBlocks;
    *pBlocksAddress = (IndexType *)blocksInfo;
}

int get_distribution(int dimension, LoopBounds *loopBounds, int threads[3], int *pResBlocks, IndexType **pBlocksAddress) {
    static Cache **caches = NULL;
//    static int tmp = 0;
//    tmp++;
    int hit = 0;
    if (!caches) {
        int dc;
        assertCuda(cudaGetDeviceCount(&dc));
        dvmh_log(DEBUG, "Creating cache for %d devices", dc);
        caches = (Cache **)malloc(sizeof(Cache *) * dc);
        int i;
        for (i = 0; i < dc; i++)
            caches[i] = init_cache(12);
    }
    int myIndex;
    assertCuda(cudaGetDevice(&myIndex));
    Cache *cache = caches[myIndex];
#ifndef DONT_USE_CACHE
    hit = probe_cache(cache, dimension, loopBounds, threads, pResBlocks, pBlocksAddress);
    dvmh_log(TRACE, "Cache probe: %s", hit ? "HIT!" : "miss");
#endif
    if (!hit) {
        IndexType *blocksInfo;
        genBlocksInfo(dimension, loopBounds, threads, pResBlocks, &blocksInfo);
        int resBlocks = *pResBlocks;
        IndexType *deviceBlocksInfo;
        assertCuda(cudaMalloc((void **)&deviceBlocksInfo, sizeof(IndexType) * resBlocks * dimension * 2));
        assertCuda(cudaMemcpy(deviceBlocksInfo, blocksInfo, sizeof(IndexType) * resBlocks * dimension * 2, cudaMemcpyHostToDevice));
        *pBlocksAddress = deviceBlocksInfo;

        add_to_cache(cache, dimension, loopBounds, threads, resBlocks, deviceBlocksInfo);
    }
    return hit;
}
